{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cfb64210-9c6b-47d7-81f4-67dbdab68e4c",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Composable Indices Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa0e62b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e27b0473-4bda-47f0-b6ed-fd482eac1a13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import (\n",
    "    GPTSimpleVectorIndex, \n",
    "    GPTSimpleKeywordTableIndex, \n",
    "    GPTListIndex, \n",
    "    SimpleDirectoryReader\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49e0d841-680f-4a0c-b455-788b54978ebf",
   "metadata": {},
   "source": [
    "#### Load Datasets\n",
    "\n",
    "Load both the NYC Wikipedia page as well as Paul Graham's \"What I Worked On\" essay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ec16a8b-6aae-4bf7-9b83-b82087b4ea52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fetch \"New York City\" page from Wikipedia\n",
    "from pathlib import Path\n",
    "\n",
    "import requests\n",
    "response = requests.get(\n",
    "    'https://en.wikipedia.org/w/api.php',\n",
    "    params={\n",
    "        'action': 'query',\n",
    "        'format': 'json',\n",
    "        'titles': 'New York City',\n",
    "        'prop': 'extracts',\n",
    "        # 'exintro': True,\n",
    "        'explaintext': True,\n",
    "    }\n",
    ").json()\n",
    "page = next(iter(response['query']['pages'].values()))\n",
    "nyc_text = page['extract']\n",
    "\n",
    "data_path = Path('data')\n",
    "if not data_path.exists():\n",
    "    Path.mkdir(data_path)\n",
    "\n",
    "with open('../test_wiki/data/nyc_text.txt', 'w') as fp:\n",
    "    fp.write(nyc_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c00aeb-adef-4ce3-8134-031de18e64ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load NYC dataset\n",
    "nyc_documents = SimpleDirectoryReader('../test_wiki/data/').load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddff8f98-e002-40c5-93ac-93aa40dca5ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load PG's essay\n",
    "essay_documents = SimpleDirectoryReader('../paul_graham_essay/data/').load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1782198-c0de-4679-8951-1297c21b8639",
   "metadata": {},
   "source": [
    "### Building the document indices\n",
    "Build a tree index for the NYC wiki page and PG essay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5431e83e-428b-4473-bad1-24b7a6c4db38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# build NYC index\n",
    "nyc_index = GPTSimpleVectorIndex.from_documents(nyc_documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ebb00ff-e9c0-4ec9-ac05-43d4ab7a0d0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "nyc_index.save_to_disk('index_nyc.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b5aad4a-49ef-4b24-962a-0793f4f09316",
   "metadata": {},
   "outputs": [],
   "source": [
    "# build essay index\n",
    "essay_index = GPTSimpleVectorIndex.from_documents(essay_documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de0d248b-bc87-4129-a12f-c96eda46f72e",
   "metadata": {},
   "outputs": [],
   "source": [
    "essay_index.save_to_disk('index_pg.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8aaf556-df77-4fac-812b-0b6c6d1da0ef",
   "metadata": {},
   "source": [
    "### Loading the indices\n",
    "Build a tree indices for the NYC wiki page and PG essay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98068ef8-aead-46e7-8dac-0d05b5a86e6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try loading\n",
    "nyc_index = GPTSimpleVectorIndex.load_from_disk('index_nyc.json')\n",
    "essay_index = GPTSimpleVectorIndex.load_from_disk('index_pg.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdcb22d5-4df8-4d65-aa29-6493fc027fe2",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Set summaries for the indices\n",
    "\n",
    "Add text summaries to indices, so we can compose other indices on top of it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4149cbbd-7d0b-48c4-8c47-7d67ae0c55f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "nyc_index_summary = \"\"\"\n",
    "    New York, often called New York City or NYC, \n",
    "    is the most populous city in the United States. \n",
    "    With a 2020 population of 8,804,190 distributed over 300.46 square miles (778.2 km2), \n",
    "    New York City is also the most densely populated major city in the United States, \n",
    "    and is more than twice as populous as second-place Los Angeles. \n",
    "    New York City lies at the southern tip of New York State, and \n",
    "    constitutes the geographical and demographic center of both the \n",
    "    Northeast megalopolis and the New York metropolitan area, the \n",
    "    largest metropolitan area in the world by urban landmass.[8] With over \n",
    "    20.1 million people in its metropolitan statistical area and 23.5 million \n",
    "    in its combined statistical area as of 2020, New York is one of the world's \n",
    "    most populous megacities, and over 58 million people live within 250 mi (400 km) of \n",
    "    the city. New York City is a global cultural, financial, and media center with \n",
    "    a significant influence on commerce, health care and life sciences, entertainment, \n",
    "    research, technology, education, politics, tourism, dining, art, fashion, and sports. \n",
    "    Home to the headquarters of the United Nations, \n",
    "    New York is an important center for international diplomacy,\n",
    "    an established safe haven for global investors, and is sometimes described as the capital of the world.\n",
    "\"\"\"\n",
    "essay_index_summary = \"\"\"\n",
    "    Author: Paul Graham. \n",
    "    The author grew up painting and writing essays. \n",
    "    He wrote a book on Lisp and did freelance Lisp hacking work to support himself. \n",
    "    He also became the de facto studio assistant for Idelle Weber, an early photorealist painter. \n",
    "    He eventually had the idea to start a company to put art galleries online, but the idea was unsuccessful. \n",
    "    He then had the idea to write software to build online stores, which became the basis for his successful company, Viaweb. \n",
    "    After Viaweb was acquired by Yahoo!, the author returned to painting and started writing essays online. \n",
    "    He wrote a book of essays, Hackers & Painters, and worked on spam filters. \n",
    "    He also bought a building in Cambridge to use as an office. \n",
    "    He then had the idea to start Y Combinator, an investment firm that would \n",
    "    make a larger number of smaller investments and help founders remain as CEO. \n",
    "    He and his partner Jessica Livingston ran Y Combinator and funded a batch of startups twice a year. \n",
    "    He also continued to write essays, cook for groups of friends, and explore the concept of invented vs discovered in software. \n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4d3cd8b-4134-4cfa-8002-e0a34694d2e1",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Build Keyword Table Index on top of tree indices! \n",
    "\n",
    "We set summaries for each of the NYC and essay indices, and then compose a keyword index on top of it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d68750c-e5ae-481a-8b03-6173020c9bf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.composability import ComposableGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76c251ca-b06b-42e9-ac99-aa0a0a5187d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set query config\n",
    "query_configs = [\n",
    "    {\n",
    "        \"index_struct_type\": \"simple_dict\",\n",
    "        \"query_mode\": \"default\",\n",
    "        \"query_kwargs\": {\n",
    "            \"similarity_top_k\": 1\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"index_struct_type\": \"keyword_table\",\n",
    "        \"query_mode\": \"simple\",\n",
    "        \"query_kwargs\": {}\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f975514f-fddd-4737-91de-97bc61394ea9",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = ComposableGraph.from_indices(\n",
    "    GPTSimpleKeywordTableIndex,\n",
    "    [nyc_index, essay_index], \n",
    "    index_summaries=[nyc_index_summary, essay_index_summary],\n",
    "    max_keywords_per_chunk=50\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae127943-afac-48b4-b22d-84a37e553e4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [optional] save to disk\n",
    "graph.save_to_disk(\"index_graph.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dca2b64b-9af1-456f-8dab-822bfdc5d0ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [optional] load from disk\n",
    "graph = ComposableGraph.load_from_disk(\"index_graph.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c4e58b-b153-4e43-bc02-274a85babbe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set Logging to DEBUG for more detailed outputs\n",
    "# ask it a question about NYC \n",
    "response = graph.query(\n",
    "    \"What is the climate of New York City like? How cold is it during the winter?\", \n",
    "    query_configs=query_configs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0a43443-3e00-4e48-b3ab-f6369191d53a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c78bc3da-6bad-4998-9a81-90a3fa9200a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get source of response\n",
    "print(response.get_formatted_sources())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b53e45e-93aa-4b49-a497-ab403f6254f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ask it a question about PG's essay\n",
    "response = graph.query(\n",
    "    \"What did the author do growing up, before his time at Y Combinator?\", \n",
    "    query_configs=query_configs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06dc71bb-882d-49f5-8566-69b0ea5019dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0894565-2b2c-4987-a891-17ba44d775b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get source of response\n",
    "print(response.get_formatted_sources())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce7efacd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
